﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Net.Security;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Security;
using System.Security.Authentication;
using System.Security.Cryptography.X509Certificates;
using System.Threading;
using xxpl1.Log;
using xxpl1.Utils;

namespace xxpl1.Cert
{
    //1. makecert -r -pe -n "CN=Test Certificate" -sky exchange -sv testcert.pvk testcert.cer
    //2. pvk2pfx -pvk testcert.pvk -spc testcert.cer -pfx testcert.pfx
    public static class SSPICertUtils
    {
        private static readonly ILog logger = LogManager.GetLogger(typeof(SSPICertUtils));
        private static readonly string certFolder = Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "cert/");
        private static readonly Dictionary<string, X509Certificate> certCache = new Dictionary<string, X509Certificate>();

        public static SslStream WrapStream(string host, Stream networkStream)
        {
            X509Certificate certificate = GetCert(host);

            var sslStream = new SslStream(networkStream, false);

            sslStream.AuthenticateAsServer(certificate, false, SslProtocols.Tls, false);

            return sslStream;
        }

        public static X509Certificate GetCert(string host)
        {
            if (!File.Exists(Path.Combine(certFolder, "_ca.pfx")))
            {
                MakeCA();
            }

            if (host.Contains(":"))
                host = host.Substring(0, host.IndexOf(":"));

            string certPath = Path.Combine(certFolder, host + ".pfx");

            if (!certCache.ContainsKey(certPath))
            {
                lock (certCache)
                {
                    if (!certCache.ContainsKey(certPath))
                    {
                        if (!File.Exists(certPath))
                        {
                            MakeCert(host);
                        }
                        certCache.Add(certPath, new X509Certificate2(certPath));
                    }
                }
            }

            return certCache[certPath];
        }

        private static void MakeCert(string host)
        {
            string capvk = Path.Combine(certFolder, "_ca.pvk");
            string cacer = Path.Combine(certFolder, "_ca.cer");
            string pvk = Path.Combine(certFolder, host + ".pvk");
            string cer = Path.Combine(certFolder, host + ".cer");
            string pfx = Path.Combine(certFolder, host + ".pfx");
            var makeCertProcessInfo = new ProcessStartInfo("makecert.exe")
                                          {
                                              Arguments = string.Format("-iv \"{0}\" -ic \"{1}\" -pe -n \"CN={2}\" -sky exchange -sv \"{3}\" \"{4}\"", capvk, cacer, host, pvk, cer),
                                              UseShellExecute = false,
                                              CreateNoWindow = true,
                                              RedirectStandardOutput = true
                                          };

            Process makeCertProcess = Process.Start(makeCertProcessInfo);

            //SetNoPass();

            logger.Debug(makeCertProcess.StandardOutput.ReadToEnd());

            makeCertProcess.WaitForExit();

            var ps2 = new ProcessStartInfo("pvk2pfx.exe")
                          {
                              Arguments = string.Format("-pvk \"{0}\" -spc \"{1}\" -pfx \"{2}\"", pvk, cer, pfx),
                              UseShellExecute = false,
                              CreateNoWindow = true,
                              RedirectStandardOutput = true
                          };

            Process p2 = Process.Start(ps2);

            logger.Debug(p2.StandardOutput.ReadToEnd());

            p2.WaitForExit();
        }

        private static void MakeCA()
        {
            string capvk = Path.Combine(certFolder, "_ca.pvk");
            string cacer = Path.Combine(certFolder, "_ca.cer");
            string capfx = Path.Combine(certFolder, "_ca.pfx");
            var makeCertProcessInfo = new ProcessStartInfo("makecert.exe")
                                          {
                                              Arguments = string.Format("-r -pe -n \"CN=SSPI Proxy Root Certificate\" -sky exchange -sv \"{0}\" \"{1}\"", capvk, cacer),
                                              UseShellExecute = false,
                                              CreateNoWindow = true,
                                              RedirectStandardOutput = true
                                          };

            Process p = Process.Start(makeCertProcessInfo);

            //SetNoPass();

            logger.Debug(p.StandardOutput.ReadToEnd());

            p.WaitForExit();

            var ps2 = new ProcessStartInfo("pvk2pfx.exe")
                          {
                              Arguments = string.Format("-pvk \"{0}\" -spc \"{1}\" -pfx \"{2}\"", capvk, cacer, capfx),
                              UseShellExecute = false,
                              CreateNoWindow = true,
                              RedirectStandardOutput = true
                          };

            Process p2 = Process.Start(ps2);

            logger.Debug(p2.StandardOutput.ReadToEnd());

            p2.WaitForExit();
        }

        public static void SetNoPass()
        {
            int i = 0;
            while (i++ < 500)
            {
                IntPtr intPtr = Win32.FindWindow(null, "Create Private Key Password");
                if (intPtr != IntPtr.Zero)
                {
                    intPtr = Win32.FindWindowEx(intPtr, IntPtr.Zero, "Button", "None");
                    Win32.SendMessage(intPtr, (uint)Win32.ButtonControlMessages.BM_CLICK, IntPtr.Zero, IntPtr.Zero);
                    break;
                }
                Thread.Sleep(50);
            }
        }

        public static byte[] CreateSelfSignCertificatePfx(
            string x500,
            DateTime startTime,
            DateTime endTime)
        {
            byte[] pfxData = CreateSelfSignCertificatePfx(
                x500,
                startTime,
                endTime,
                (SecureString)null);
            return pfxData;
        }

        public static byte[] CreateSelfSignCertificatePfx(
            string x500,
            DateTime startTime,
            DateTime endTime,
            string insecurePassword)
        {
            byte[] pfxData;
            SecureString password = null;

            try
            {
                if (!string.IsNullOrEmpty(insecurePassword))
                {
                    password = new SecureString();
                    foreach (char ch in insecurePassword)
                    {
                        password.AppendChar(ch);
                    }

                    password.MakeReadOnly();
                }

                pfxData = CreateSelfSignCertificatePfx(
                    x500,
                    startTime,
                    endTime,
                    password);
            }
            finally
            {
                if (password != null)
                {
                    password.Dispose();
                }
            }

            return pfxData;
        }

        public static byte[] CreateSelfSignCertificatePfx(
            string x500,
            DateTime startTime,
            DateTime endTime,
            SecureString password)
        {
            byte[] pfxData;

            if (x500 == null)
            {
                x500 = "";
            }

            SystemTime startSystemTime = ToSystemTime(startTime);
            SystemTime endSystemTime = ToSystemTime(endTime);
            string containerName = Guid.NewGuid().ToString();

            var dataHandle = new GCHandle();
            IntPtr providerContext = IntPtr.Zero;
            IntPtr cryptKey = IntPtr.Zero;
            IntPtr certContext = IntPtr.Zero;
            IntPtr certStore = IntPtr.Zero;
            IntPtr storeCertContext = IntPtr.Zero;
            IntPtr passwordPtr = IntPtr.Zero;
            RuntimeHelpers.PrepareConstrainedRegions();
            try
            {
                Check(NativeMethods.CryptAcquireContextW(
                    out providerContext,
                    containerName,
                    null,
                    1, // PROV_RSA_FULL
                    8)); // CRYPT_NEWKEYSET

                Check(NativeMethods.CryptGenKey(
                    providerContext,
                    1, // AT_KEYEXCHANGE
                    1, // CRYPT_EXPORTABLE
                    out cryptKey));

                IntPtr errorStringPtr;
                int nameDataLength = 0;

                // errorStringPtr gets a pointer into the middle of the x500 string,
                // so x500 needs to be pinned until after we've copied the value
                // of errorStringPtr.
                dataHandle = GCHandle.Alloc(x500, GCHandleType.Pinned);

                if (!NativeMethods.CertStrToNameW(
                    0x00010001, // X509_ASN_ENCODING | PKCS_7_ASN_ENCODING
                    dataHandle.AddrOfPinnedObject(),
                    3, // CERT_X500_NAME_STR = 3
                    IntPtr.Zero,
                    null,
                    ref nameDataLength,
                    out errorStringPtr))
                {
                    string error = Marshal.PtrToStringUni(errorStringPtr);
                    throw new ArgumentException(error);
                }

                var nameData = new byte[nameDataLength];

                if (!NativeMethods.CertStrToNameW(
                    0x00010001, // X509_ASN_ENCODING | PKCS_7_ASN_ENCODING
                    dataHandle.AddrOfPinnedObject(),
                    3, // CERT_X500_NAME_STR = 3
                    IntPtr.Zero,
                    nameData,
                    ref nameDataLength,
                    out errorStringPtr))
                {
                    string error = Marshal.PtrToStringUni(errorStringPtr);
                    throw new ArgumentException(error);
                }

                dataHandle.Free();

                dataHandle = GCHandle.Alloc(nameData, GCHandleType.Pinned);
                var nameBlob = new CryptoApiBlob(
                    nameData.Length,
                    dataHandle.AddrOfPinnedObject());

                var kpi = new CryptKeyProviderInformation();
                kpi.ContainerName = containerName;
                kpi.ProviderType = 1; // PROV_RSA_FULL
                kpi.KeySpec = 1; // AT_KEYEXCHANGE

                certContext = NativeMethods.CertCreateSelfSignCertificate(
                    providerContext,
                    ref nameBlob,
                    0,
                    ref kpi,
                    IntPtr.Zero, // default = SHA1RSA
                    ref startSystemTime,
                    ref endSystemTime,
                    IntPtr.Zero);
                Check(certContext != IntPtr.Zero);
                dataHandle.Free();

                certStore = NativeMethods.CertOpenStore(
                    "Memory", // sz_CERT_STORE_PROV_MEMORY
                    0,
                    IntPtr.Zero,
                    0x2000, // CERT_STORE_CREATE_NEW_FLAG
                    IntPtr.Zero);
                Check(certStore != IntPtr.Zero);

                Check(NativeMethods.CertAddCertificateContextToStore(
                    certStore,
                    certContext,
                    1, // CERT_STORE_ADD_NEW
                    out storeCertContext));

                NativeMethods.CertSetCertificateContextProperty(
                    storeCertContext,
                    2, // CERT_KEY_PROV_INFO_PROP_ID
                    0,
                    ref kpi);

                if (password != null)
                {
                    passwordPtr = Marshal.SecureStringToCoTaskMemUnicode(password);
                }

                var pfxBlob = new CryptoApiBlob();
                Check(NativeMethods.PFXExportCertStoreEx(
                    certStore,
                    ref pfxBlob,
                    passwordPtr,
                    IntPtr.Zero,
                    7)); // EXPORT_PRIVATE_KEYS | REPORT_NO_PRIVATE_KEY | REPORT_NOT_ABLE_TO_EXPORT_PRIVATE_KEY

                pfxData = new byte[pfxBlob.DataLength];
                dataHandle = GCHandle.Alloc(pfxData, GCHandleType.Pinned);
                pfxBlob.Data = dataHandle.AddrOfPinnedObject();
                Check(NativeMethods.PFXExportCertStoreEx(
                    certStore,
                    ref pfxBlob,
                    passwordPtr,
                    IntPtr.Zero,
                    7)); // EXPORT_PRIVATE_KEYS | REPORT_NO_PRIVATE_KEY | REPORT_NOT_ABLE_TO_EXPORT_PRIVATE_KEY
                dataHandle.Free();
            }
            finally
            {
                if (passwordPtr != IntPtr.Zero)
                {
                    Marshal.ZeroFreeCoTaskMemUnicode(passwordPtr);
                }

                if (dataHandle.IsAllocated)
                {
                    dataHandle.Free();
                }

                if (certContext != IntPtr.Zero)
                {
                    NativeMethods.CertFreeCertificateContext(certContext);
                }

                if (storeCertContext != IntPtr.Zero)
                {
                    NativeMethods.CertFreeCertificateContext(storeCertContext);
                }

                if (certStore != IntPtr.Zero)
                {
                    NativeMethods.CertCloseStore(certStore, 0);
                }

                if (cryptKey != IntPtr.Zero)
                {
                    NativeMethods.CryptDestroyKey(cryptKey);
                }

                if (providerContext != IntPtr.Zero)
                {
                    NativeMethods.CryptReleaseContext(providerContext, 0);
                    NativeMethods.CryptAcquireContextW(
                        out providerContext,
                        containerName,
                        null,
                        1, // PROV_RSA_FULL
                        0x10); // CRYPT_DELETEKEYSET
                }
            }

            return pfxData;
        }

        private static SystemTime ToSystemTime(DateTime dateTime)
        {
            long fileTime = dateTime.ToFileTime();
            SystemTime systemTime;
            Check(NativeMethods.FileTimeToSystemTime(ref fileTime, out systemTime));
            return systemTime;
        }

        private static void Check(bool nativeCallSucceeded)
        {
            if (!nativeCallSucceeded)
            {
                int error = Marshal.GetHRForLastWin32Error();
                Marshal.ThrowExceptionForHR(error);
            }
        }

        #region Nested type: CryptKeyProviderInformation

        [StructLayout(LayoutKind.Sequential)]
        private struct CryptKeyProviderInformation
        {
            [MarshalAs(UnmanagedType.LPWStr)]
            public string ContainerName;
            [MarshalAs(UnmanagedType.LPWStr)]
            public readonly string ProviderName;
            public int ProviderType;
            public readonly int Flags;
            public readonly int ProviderParameterCount;
            public readonly IntPtr ProviderParameters; // PCRYPT_KEY_PROV_PARAM
            public int KeySpec;
        }

        #endregion

        #region Nested type: CryptoApiBlob

        [StructLayout(LayoutKind.Sequential)]
        private struct CryptoApiBlob
        {
            public readonly int DataLength;
            public IntPtr Data;

            public CryptoApiBlob(int dataLength, IntPtr data)
            {
                DataLength = dataLength;
                Data = data;
            }
        }

        #endregion

        #region Nested type: NativeMethods

        private static class NativeMethods
        {
            [DllImport("kernel32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool FileTimeToSystemTime(
                [In] ref long fileTime,
                out SystemTime systemTime);

            [DllImport("AdvApi32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CryptAcquireContextW(
                out IntPtr providerContext,
                [MarshalAs(UnmanagedType.LPWStr)] string container,
                [MarshalAs(UnmanagedType.LPWStr)] string provider,
                int providerType,
                int flags);

            [DllImport("AdvApi32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CryptReleaseContext(
                IntPtr providerContext,
                int flags);

            [DllImport("AdvApi32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CryptGenKey(
                IntPtr providerContext,
                int algorithmId,
                int flags,
                out IntPtr cryptKeyHandle);

            [DllImport("AdvApi32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CryptDestroyKey(
                IntPtr cryptKeyHandle);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CertStrToNameW(
                int certificateEncodingType,
                IntPtr x500,
                int strType,
                IntPtr reserved,
                [MarshalAs(UnmanagedType.LPArray)] [Out] byte[] encoded,
                ref int encodedLength,
                out IntPtr errorString);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            public static extern IntPtr CertCreateSelfSignCertificate(
                IntPtr providerHandle,
                [In] ref CryptoApiBlob subjectIssuerBlob,
                int flags,
                [In] ref CryptKeyProviderInformation keyProviderInformation,
                IntPtr signatureAlgorithm,
                [In] ref SystemTime startTime,
                [In] ref SystemTime endTime,
                IntPtr extensions);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CertFreeCertificateContext(
                IntPtr certificateContext);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            public static extern IntPtr CertOpenStore(
                [MarshalAs(UnmanagedType.LPStr)] string storeProvider,
                int messageAndCertificateEncodingType,
                IntPtr cryptProvHandle,
                int flags,
                IntPtr parameters);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CertCloseStore(
                IntPtr certificateStoreHandle,
                int flags);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CertAddCertificateContextToStore(
                IntPtr certificateStoreHandle,
                IntPtr certificateContext,
                int addDisposition,
                out IntPtr storeContextPtr);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool CertSetCertificateContextProperty(
                IntPtr certificateContext,
                int propertyId,
                int flags,
                [In] ref CryptKeyProviderInformation data);

            [DllImport("Crypt32.dll", SetLastError = true, ExactSpelling = true)]
            [return: MarshalAs(UnmanagedType.Bool)]
            public static extern bool PFXExportCertStoreEx(
                IntPtr certificateStoreHandle,
                ref CryptoApiBlob pfxBlob,
                IntPtr password,
                IntPtr reserved,
                int flags);
        }

        #endregion

        #region Nested type: SystemTime

        [StructLayout(LayoutKind.Sequential)]
        private struct SystemTime
        {
            public readonly short Year;
            public readonly short Month;
            public readonly short DayOfWeek;
            public readonly short Day;
            public readonly short Hour;
            public readonly short Minute;
            public readonly short Second;
            public readonly short Milliseconds;
        }

        #endregion
    }
}